import mindspore.nn as nn
import mindspore.ops as ops


class ResidualBlock(nn.Cell):
    def __init__(self, in_channels, out_channels, stride=1):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels,
                               kernel_size=3, stride=stride,
                               padding=0, pad_mode="same")
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels,
                               kernel_size=3, stride=1,
                               padding=0, pad_mode="same")
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()
        # shortcut
        self.down_sample = False
        self.down_sample_layer = None
        if in_channels != out_channels:
            self.down_sample = True
            self.down_sample_layer = nn.SequentialCell([
                nn.Conv2d(in_channels, out_channels,
                          kernel_size=1, stride=stride,
                          padding=0, pad_mode="same"),
                nn.BatchNorm2d(out_channels)
            ])

    def construct(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.down_sample:
            identity = self.down_sample_layer(identity)
        out = out + identity
        out = self.relu(out)

        return out


class ResNet18(nn.Cell):
    def __init__(self, num_classes=10):
        super(ResNet18, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu1 = nn.ReLU()
        self.conv2_x_max_pool = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv2_x_1 = ResidualBlock(64, 64)
        self.conv2_x_2 = ResidualBlock(64, 64)
        self.conv3_x_1 = ResidualBlock(64, 128)
        self.conv3_x_2 = ResidualBlock(128, 128)
        self.conv4_x_1 = ResidualBlock(128, 256)
        self.conv4_x_2 = ResidualBlock(256, 256)
        self.conv5_x_1 = ResidualBlock(256, 512)
        self.conv5_x_2 = ResidualBlock(512, 512)
        self.mean = ops.ReduceMean(keep_dims=True)
        self.flatten = nn.Flatten()
        self.dense = nn.Dense(512, num_classes, has_bias=True)

    def construct(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu1(out)
        # layer 1
        out = self.conv2_x_max_pool(out)
        out = self.conv2_x_1(out)
        out = self.conv2_x_2(out)
        # layer 2
        out = self.conv3_x_1(out)
        out = self.conv3_x_2(out)
        # layer 3
        out = self.conv4_x_1(out)
        out = self.conv4_x_2(out)
        # layer 4
        out = self.conv5_x_1(out)
        out = self.conv5_x_2(out)
        # header
        out = self.mean(out, (2, 3))
        out = self.flatten(out)
        out = self.dense(out)

        return out
